package com.xzzz.expand.es.core;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.StrUtil;
import com.xzzz.common.base.exception.XzException500;
import com.xzzz.common.base.util.json.JsonUtil;
import com.xzzz.expand.es.core.annotation.EsId;
import com.xzzz.expand.es.core.annotation.EsIndex;
import com.xzzz.expand.es.core.utils.EsUtil;
import lombok.NonNull;
import lombok.extern.slf4j.Slf4j;
import org.elasticsearch.ElasticsearchStatusException;
import org.elasticsearch.action.ActionListener;
import org.elasticsearch.action.bulk.BulkRequest;
import org.elasticsearch.action.bulk.BulkResponse;
import org.elasticsearch.action.delete.DeleteRequest;
import org.elasticsearch.action.delete.DeleteResponse;
import org.elasticsearch.action.get.GetRequest;
import org.elasticsearch.action.get.GetResponse;
import org.elasticsearch.action.index.IndexRequest;
import org.elasticsearch.action.index.IndexResponse;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.action.support.WriteRequest;
import org.elasticsearch.action.update.UpdateRequest;
import org.elasticsearch.action.update.UpdateResponse;
import org.elasticsearch.client.Cancellable;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.client.core.CountRequest;
import org.elasticsearch.client.core.CountResponse;
import org.elasticsearch.common.xcontent.XContentType;
import org.elasticsearch.index.query.QueryBuilder;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.index.reindex.BulkByScrollResponse;
import org.elasticsearch.index.reindex.DeleteByQueryRequest;
import org.elasticsearch.rest.RestStatus;
import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.SearchHits;
import org.elasticsearch.search.builder.SearchSourceBuilder;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;

/**
 * ES 增删改查
 * <ul>
 *  查询
 *  <ol>
 *      <li>判断是否存在 {@link EsCurdManager#exists(String, String)}</li>
 *      <li>查询数量 {@link EsCurdManager#count(String, QueryBuilder)}</li>
 *      <li>根据ID查询 {@link EsCurdManager#get(String, String, Class)}</li>
 *      <li>根据ID集合查询 {@link EsCurdManager#get(String, List, Class)}</li>
 *  </ol>
 * <p>
 * <p>
 *  新增
 *  <ol>
 *      <li>单个新增 {@link EsCurdManager#add(Object)}</li>
 *      <li>单个新增 {@link EsCurdManager#add(String, Object)}</li>
 *      <li>单个新增 {@link EsCurdManager#add(String, Object, boolean)}</li>
 *      <li>批量新增 {@link EsCurdManager#addBatch(List)}</li>
 *      <li>批量新增 {@link EsCurdManager#addBatch(String, List)}</li>
 *      <li>批量新增 {@link EsCurdManager#addBatch(String, List, boolean)}</li>
 *  </ol>
 * <p>
 * <p>
 *  修改
 *  <ol>
 *      <li>单个修改 {@link EsCurdManager#update(Object, boolean)}</li>
 *      <li>批量修改 {@link EsCurdManager#updateBatch(List, boolean)}</li>
 *  </ol>
 * <p>
 * <p>
 *  删除
 *  <ol>
 *
 *  </ol>
 * </ul>
 */
@Slf4j
public class EsCurdManager extends EsSearchManager {

    public EsCurdManager(RestHighLevelClient restHighLevelClient, EsProperties esProperties) {
        super(restHighLevelClient, esProperties);
    }

    // region -----------------------------------------------< 查询数据 >------------------------------------------------

    /**
     * 判断是否存在
     *
     * @param index index
     * @param id    id
     * @return 是否存在
     */
    public boolean exists(@NonNull String index, @NonNull String id) {
        try {
            GetRequest request = new GetRequest(index, id);
            return client.exists(request, RequestOptions.DEFAULT);
        } catch (ElasticsearchStatusException ex) {
            esExceptionSuppose(ex);
            throw new XzException500("ES处理错误:" + ex.getMessage());
        } catch (IOException e) {
            throw new XzException500("查询数据失败:" + e.getMessage());
        }
    }

    /**
     * 查询条数
     *
     * @param index        index
     * @param queryBuilder 查询条件, 如果不传入查询条件, 则会查该索引下的文档数量
     * @return 查询条数
     */
    public long count(@NonNull String index, QueryBuilder queryBuilder) {
        CountRequest request = new CountRequest(index);
        if (queryBuilder != null) {
            request.query(queryBuilder);
        }
        try {
            CountResponse countResponse = client.count(request, RequestOptions.DEFAULT);
            return countResponse.getCount();
        } catch (ElasticsearchStatusException ex) {
            esExceptionSuppose(ex);
            throw new XzException500("ES处理错误:" + ex.getMessage());
        } catch (IOException e) {
            throw new XzException500("查询数据失败:" + e.getMessage());
        }
    }

    /**
     * 查询文档原文
     *
     * @param index 索引
     * @param id    id
     * @return 索引原对象
     */
    public Map<String, Object> getSource(@NonNull String index, @NonNull String id) {
        GetRequest request = new GetRequest(index, id);
        try {
            GetResponse response = client.get(request, RequestOptions.DEFAULT);
            if (!response.isSourceEmpty() && response.isExists()) {
                return response.getSource();
            }

            return null;
        } catch (ElasticsearchStatusException ex) {
            esExceptionSuppose(ex);
            throw new XzException500("ES处理错误:" + ex.getMessage());
        } catch (IOException e) {
            throw new XzException500("查询数据失败:" + e.getMessage());
        }
    }

    /**
     * 根据 id 查询文档
     *
     * @param index 索引
     * @param id    id
     * @param clazz 响应结果类
     * @param <T>   响应结果类类型
     * @return 响应结果
     */
    public <T> T get(@NonNull String index, @NonNull String id, Class<T> clazz) {
        Map<String, Object> source = getSource(index, id);
        if (source != null) {
            return JsonUtil.toObj(source, clazz);
        }
        return null;
    }

    /**
     * 根据 ids 查询文档<br/>
     *
     * <p>如需要更多条件, 如指定最大查询条数, 筛选指定字段等等, 请使用 {@link com.xzzz.expand.es.core.EsSearchManager#search(String, SearchSourceBuilder)}
     * 并通过 {@link QueryBuilders#idsQuery()} 查询
     *
     * <p>更多使用方法可参阅:<a href="https://www.yuque.com/xiaozeizeizi/learning/cawmb8">《API/RestHighLevelClient 速查表》</a>
     *
     * @param index 索引名称
     * @param ids   id集合
     * @param clazz 响应结果类
     * @param <T>   响应结果类类型
     * @return 响应结果
     */
    public <T> List<T> get(@NonNull String index, @NonNull List<String> ids, Class<T> clazz) {
        SearchSourceBuilder searchSource = new SearchSourceBuilder()
                .from(0)
                .size(ids.size())
                .query(QueryBuilders.idsQuery().addIds(ids.toArray(new String[]{})));
        SearchResponse response = this.search(index, searchSource);

        List<T> resultList = new ArrayList<>();

        SearchHits hits = response.getHits();
        if (hits.getTotalHits().value != 0) {
            for (SearchHit hit : hits) {
                if (!hit.hasSource()) {
                    continue;
                }
                T result = JsonUtil.toObj(hit.getSourceAsMap(), clazz);
                resultList.add(result);
            }
        }

        return resultList;
    }

    // endregion


    // region -----------------------------------------------< 新增数据 >------------------------------------------------

    /**
     * 新增数据, ID存在时会报错, 如需要修改请使用修改功能 {@link #update(Object, boolean)}
     *
     * @param index   索引名称, 如果传入该值, 将忽略对象中的 {@link EsIndex} 注解
     * @param source  数据对象
     * @param waitFor 刷新策略, 为 true 则等待此变动可查询时才返回
     */
    public void add(String index, @NonNull Object source, boolean waitFor) {
        if (StrUtil.isBlank(index)) {
            index = EsUtil.getIndexByAnnotation(source);
        }
        String id = EsUtil.getIdByAnnotation(source);
        try {
            IndexRequest request = new IndexRequest(index)
                    .id(id)
                    .source(JsonUtil.toJson(source), XContentType.JSON);
            if (waitFor) {
                request.setRefreshPolicy(WriteRequest.RefreshPolicy.WAIT_UNTIL);
            }
            IndexResponse response = client.index(request, RequestOptions.DEFAULT);
            /*
             * 200: 成功
             * 201: 创建数据成功
             */
            if (response.status() == RestStatus.OK ||
                    response.status() == RestStatus.CREATED) {
                log.debug("添加数据成功 索引:{}, response 状态:{}, id:{}", index, response.status().getStatus(), response.getId());
                return;
            }
            throw new XzException500("保存失败");
        } catch (ElasticsearchStatusException ex) {
            esExceptionSuppose(ex);
            throw new RuntimeException("ES处理错误:" + ex.getMessage());
        } catch (IOException e) {
            e.printStackTrace();
            throw new RuntimeException(String.format("新增索引错误: %s, 对象: %s", e.getMessage(), source.toString()));
        }
    }

    /**
     * 新增数据
     *
     * @param index  索引名称
     * @param source 数据对象
     */
    public void add(String index, @NonNull Object source) {
        this.add(index, source, false);
    }

    /**
     * 新增数据
     *
     * @param source 数据对象
     */
    public void add(@NonNull Object source) {
        this.add(null, source, false);
    }

    /**
     * 批量新增
     * <p>1. 需要保证 sources 中的对象全部具有相同的 indexName, 只会使用下标 [0] 的 indexName
     * <p>2. 如果对象中不包含 {@link EsId} 所注释的字段, 或该字段的值为 null 或 "", 将会自动生成一个ID
     *
     * @param index   索引名称, 如果传入该值, 将忽略对象中的 {@link EsIndex} 注解
     * @param sources 数据对象集合
     * @param waitFor 刷新策略, 为 true 则等待此变动可查询时才返回
     */
    public void addBatch(String index, @NonNull List<?> sources, boolean waitFor) {
        if (CollUtil.isEmpty(sources)) {
            return;
        }
        BulkRequest request = new BulkRequest();
        if (StrUtil.isBlank(index)) {
            index = EsUtil.getIndexByAnnotation(sources.get(0));
        }
        for (Object source : sources) {
            IndexRequest indexRequest = new IndexRequest()
                    .index(index)
                    .id(EsUtil.getIdByAnnotation(source))
                    .source(JsonUtil.toJson(source), XContentType.JSON);
            request.add(indexRequest);
        }
        batch(request, waitFor);
    }

    /**
     * 指定 indexName, 将会忽略 sources 中的注解
     *
     * @param index   indexName
     * @param sources 数据对象集合
     */
    public void addBatch(String index, @NonNull List<?> sources) {
        this.addBatch(index, sources, false);
    }

    /**
     * 批量新增数据
     *
     * @param sources 数据对象集合
     */
    public void addBatch(@NonNull List<?> sources) {
        this.addBatch(null, sources, false);
    }

    // endregion


    // region -----------------------------------------------< 修改数据 >------------------------------------------------

    /**
     * 修改数据
     *
     * @param source  数据对象
     * @param waitFor 刷新策略, 为 true 则等待此变动可查询时才返回
     */
    public void update(@NonNull Object source, boolean waitFor) {
        String indexName = EsUtil.getIndexByAnnotation(source);
        String id = EsUtil.getIdByAnnotation(source);

        try {
            UpdateRequest request = new UpdateRequest(indexName, id).doc(JsonUtil.toJson(source), XContentType.JSON);
            if (waitFor) {
                request.setRefreshPolicy(WriteRequest.RefreshPolicy.WAIT_UNTIL);
            }
            UpdateResponse response = client.update(request, RequestOptions.DEFAULT);
            if (response.status() != RestStatus.OK) {
                throw new RuntimeException("修改失败");
            }
            log.debug("修改数据成功 索引:{}, response 状态:{}, id:{}", indexName, response.status().getStatus(), response.getId());
        } catch (IOException e) {
            e.printStackTrace();
            throw new RuntimeException(String.format("修改对象错误: %s, 对象: %s", e.getMessage(), source.toString()));
        }
    }

    /**
     * 修改数据
     *
     * @param source 数据对象
     */
    public void update(@NonNull Object source) {
        this.update(source, false);
    }

    /**
     * 批量修改数据
     *
     * @param sources 数据对象集合
     * @param waitFor 刷新策略, 为 true 则等待此变动可查询时才返回
     */
    public void updateBatch(@NonNull List<?> sources, boolean waitFor) {
        BulkRequest request = new BulkRequest();
        for (Object source : sources) {
            UpdateRequest indexRequest = new UpdateRequest()
                    .index(EsUtil.getIndexByAnnotation(source))
                    .id(EsUtil.getIdByAnnotation(source, true))
                    .doc(JsonUtil.toJson(source), XContentType.JSON);
            request.add(indexRequest);
        }
        batch(request, waitFor);
    }

    /**
     * 批量修改数据
     *
     * @param sources 数据对象
     */
    public void updateBatch(@NonNull List<?> sources) {
        this.updateBatch(sources, false);
    }

    // endregion


    // region -----------------------------------------------< 删除数据 >------------------------------------------------

    /**
     * 删除数据
     *
     * @param index   index
     * @param id      id
     * @param waitFor 刷新策略, 为 true 则等待此变动可查询时才返回
     */
    public void delete(@NonNull String index, @NonNull String id, boolean waitFor) {
        // 执行客户端请求
        try {
            DeleteRequest request = new DeleteRequest(index, id);
            if (waitFor) {
                request.setRefreshPolicy(WriteRequest.RefreshPolicy.WAIT_UNTIL);
            }
            DeleteResponse response = client.delete(request, RequestOptions.DEFAULT);
            if (response.status() != RestStatus.OK) {
                throw new RuntimeException("删除失败");
            }
        } catch (ElasticsearchStatusException ex) {
            esExceptionSuppose(ex);
            throw new RuntimeException("ES处理错误:" + ex.getMessage());
        } catch (IOException e) {
            e.printStackTrace();
            throw new RuntimeException(String.format("删除对象错误, [INDEX:%s] [ID:%s] %s", index, id, e.getMessage()));
        }
    }

    /**
     * 删除数据
     *
     * @param index index
     * @param id    id
     */
    public void delete(@NonNull String index, @NonNull String id) {
        this.delete(index, id, false);
    }

    /**
     * 根据查询条件删除
     * <p><a href='https://www.elastic.co/guide/en/elasticsearch/reference/7.10/docs-delete-by-query.html'>Rest API 官方文档</a></p>
     * <p><a href='https://www.elastic.co/guide/en/elasticsearch/client/java-rest/7.10/java-rest-high-document-delete-by-query.html'>Java API 官方文档</a>
     *
     * @param index        索引名称
     * @param queryBuilder 查询条件
     * @return BulkByScrollResponse 删除结果
     * <p>{@link BulkByScrollResponse#getTotal()} 获取条件匹配的总条数,
     * <p>{@link BulkByScrollResponse#getDeleted()} 获取实际删除的总条数
     */
    public BulkByScrollResponse deleteByQuery(@NonNull String index, @NonNull QueryBuilder queryBuilder) {
        DeleteByQueryRequest request = new DeleteByQueryRequest(index);
        request.setQuery(queryBuilder);
        // 在请求结束后对写入的索引调用刷新
        request.setRefresh(false);
        try {
            return client.deleteByQuery(request, RequestOptions.DEFAULT);
        } catch (ElasticsearchStatusException ex) {
            esExceptionSuppose(ex);
            throw new RuntimeException("ES处理错误:" + ex.getMessage());
        } catch (IOException e) {
            e.printStackTrace();
            throw new RuntimeException(String.format("删除对象错误, [INDEX:%s] [QUERY:%s] %s", index, queryBuilder.toString(), e.getMessage()));
        }
    }

    /**
     * 根据查询条件异步删除
     *
     * <pre>{@code
     * ActionListener<BulkByScrollResponse> listener = new ActionListener<BulkByScrollResponse>() {
     *     @Override
     *     public void onResponse(BulkByScrollResponse bulkResponse) {
     *         // 删除成功
     *     }
     *     @Override
     *     public void onFailure(Exception e) {
     *         // 删除失败
     *     }
     * };}</pre>
     *
     * @param index 索引名称
     * @param queryBuilder 查询条件
     * @param actionListener 异步监听
     * @return 响应
     */
    public Cancellable deleteByQueryAsync(@NonNull String index, @NonNull QueryBuilder queryBuilder,
                                          @NonNull ActionListener<BulkByScrollResponse> actionListener) {
        DeleteByQueryRequest request = new DeleteByQueryRequest(index);
        request.setQuery(queryBuilder);
        // 在请求结束后对写入的索引调用刷新
        request.setRefresh(false);
        try {
            return client.deleteByQueryAsync(request, RequestOptions.DEFAULT, actionListener);
        } catch (ElasticsearchStatusException ex) {
            esExceptionSuppose(ex);
            throw new RuntimeException("ES处理错误:" + ex.getMessage());
        }
    }

    /**
     * 批量删除数据
     *
     * @param index   index
     * @param ids     id集合
     * @param waitFor 刷新策略, 为 true 则等待此变动可查询时才返回
     */
    public void deleteBatch(@NonNull String index, @NonNull List<String> ids, boolean waitFor) {
        BulkRequest request = new BulkRequest();
        for (String id : ids) {
            request.add(new DeleteRequest(index, id));
        }
        batch(request, waitFor);
    }

    /**
     * 批量删除数据
     *
     * @param index index
     * @param ids   id集合
     */
    public void deleteBatch(@NonNull String index, @NonNull List<String> ids) {
        this.deleteBatch(index, ids, false);
    }


    // endregion


    /**
     * 批处理请求<br/><br/>
     * <p>
     * 刷新策略 {@link WriteRequest.RefreshPolicy}
     * <ul>
     * <li>false  : 不要在此请求后刷新, 默认值. </li>
     * <li>true   : 作为请求的一部分强制刷新. 此刷新策略不能用于高索引或搜索吞吐量, 但对于具有非常低流量的索引提供一致的视图是有用的, 它非常适合测试!</li>
     * <li>waitFor: 保持此请求打开, 直到刷新使此请求的内容可以搜索. 此刷新策略与高索引和搜索吞吐量兼容, 但它会导致请求等待应答, 直到刷新发生.</li>
     * </ul>
     *
     * @param bulkRequest 批处理请求
     * @param waitFor     刷新策略 {@link WriteRequest.RefreshPolicy}, 此处只使用 waitFor, 若为true会导致请求时间变长, 批量操作
     *                    时需要酌情使用
     */
    private void batch(BulkRequest bulkRequest, boolean waitFor) {
        if (waitFor) {
            bulkRequest.setRefreshPolicy(WriteRequest.RefreshPolicy.WAIT_UNTIL);
        }
        try {
            BulkResponse response = client.bulk(bulkRequest, RequestOptions.DEFAULT);
            if (response.status() != RestStatus.OK || response.hasFailures()) {
                throw new RuntimeException(response.buildFailureMessage());
            }
        } catch (ElasticsearchStatusException ex) {
            esExceptionSuppose(ex);
            throw new RuntimeException("ES处理错误:" + ex.getMessage());
        } catch (IOException e) {
            throw new RuntimeException(String.format("批处理(%s)条错误: %S.%s",
                    bulkRequest.numberOfActions(),
                    e.getClass().getSimpleName(),
                    e.getMessage()));
        }
    }


}
